import numpy as np
import cv2
import glob
import yaml

# 定义标定板的尺寸（行数和列数）
rows = 7
cols = 7

# 准备对象点，如 (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)
objp = np.zeros((rows * cols, 3), np.float32)
objp[:, :2] = np.mgrid[0:cols, 0:rows].T.reshape(-1, 2)
objp = objp*0.25
# 存储所有图像对象点和图像点的数组
objpoints = []  # 真实世界中的3D点
imgpoints = []  # 图像中的2D点

# 读取所有图像
images = glob.glob('./cv_pkg/data/*.jpg')

for frame in images:

    img = cv2.imread(frame)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 找到标定板的角点
    ret, corners = cv2.findChessboardCorners(gray, (cols, rows), None)
    
    # 如果找到，添加对象点和图像点
    if ret == True:
        objpoints.append(objp)
        imgpoints.append(corners)

        # 绘制并显示角点
        cv2.drawChessboardCorners(img, (cols, rows), corners, ret)
        cv2.imshow('img', img)
        cv2.waitKey(500)

cv2.destroyAllWindows()

# 进行相机标定
ret, camera_matrix, dist_coeffs, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)

print("相机矩阵:\n", camera_matrix)
print("畸变系数:\n", dist_coeffs)
# print("旋转向量外参:\n", rvecs)
# print("平移向量外参:\n", tvecs)
# 保存相机参数到YAML文件
data = {
    'camera_matrix': camera_matrix.tolist(),
    'dist_coeffs': dist_coeffs.tolist()
}

with open('camera_params.yaml', 'w') as f:
    yaml.dump(data, f)

print("相机参数已保存到 camera_params.yaml")

# 计算反向投影误差
mean_error = 0
for i in range(len(objpoints)):
    imgpoints2, _ = cv2.projectPoints(objpoints[i], rvecs[i], tvecs[i], camera_matrix, dist_coeffs)
    error = cv2.norm(imgpoints[i], imgpoints2, cv2.NORM_L2) / len(imgpoints2)
    mean_error += error

print(f"总误差: {mean_error / len(objpoints)}")


# 可选：对图像进行畸变校正
for fname in images:
    img = cv2.imread(fname)
    h, w = img.shape[:2]
    new_camera_matrix, roi = cv2.getOptimalNewCameraMatrix(camera_matrix, dist_coeffs, (w, h), 1, (w, h))

    # 使用cv2.undistort进行畸变校正
    dst = cv2.undistort(img, camera_matrix, dist_coeffs, None, new_camera_matrix)

    # 裁剪图像
    x, y, w, h = roi
    dst = dst[y:y+h, x:x+w]
    cv2.imshow('img', dst)
    cv2.waitKey(500)

cv2.destroyAllWindows()